
package com.saturday.uid.worker.dao.impl;

import com.saturday.uid.config.ZookeeperProperties;
import com.saturday.uid.utils.DockerUtils;
import com.saturday.uid.utils.NetUtils;
import com.saturday.uid.worker.dao.WorkerNodeDao;
import lombok.extern.slf4j.Slf4j;
import org.I0Itec.zkclient.ZkClient;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;
import org.springframework.beans.factory.annotation.Autowired;

import java.net.SocketException;
import java.util.List;

@Slf4j
public class ZkNodeWorkerNodeDaoImpl implements WorkerNodeDao {

    @Autowired
    private ZookeeperProperties zookeeperProperties;

    @Override
    public long assignRepeatableWorkerId() {
        String zkAddress = zookeeperProperties.getAddress();
        ZkClient zkClient = new ZkClient(zkAddress,
                zookeeperProperties.getSessionTimeout(), zookeeperProperties.getConnectionTimeout());
        try {
            String repeatableUidRoot = zookeeperProperties.getRepeatableUidRoot();
            if (!zkClient.exists(repeatableUidRoot)) {
                zkClient.createPersistent(repeatableUidRoot, true);
            }

            String hostName = getAddress();
            String workNode = repeatableUidRoot + "/" + hostName;

            /**
             * 检查ip:port节点是否存在，如果存在，则取子节点的顺序编号；如果不存在，则创建持久有序节点
             */
            List<String> repeatableUidWorkNodes = zkClient.getChildren(repeatableUidRoot);
            /** 检查zk上是否存在ip:port的节点,存在，获取节点的顺序编号 */
            for (int i = 0; i < repeatableUidWorkNodes.size(); i++) {
                String workNodeResult = repeatableUidWorkNodes.get(i);
                if (workNodeResult.startsWith(hostName)) {
                    String workerIdStr = workNodeResult.substring(hostName.length());
                    long workerId = Long.valueOf(workerIdStr);
                    log.info("机器:{}, workerId复用，workerId={}", hostName, workerId);
                    return workerId;
                }
            }

            /** 不存在，创建持久有序节点, 并返回当前节点下标 */
            String workNodeResult = zkClient.create(workNode, new byte[0], CreateMode.PERSISTENT_SEQUENTIAL);
            String workerIdStr = workNodeResult.substring(workNode.length());
            long workerId = Long.valueOf(workerIdStr);
            log.info("workerId={}", workerId);
            return workerId;

        } catch (Exception e) {
            log.error("获取workerId失败:{}", e.getMessage());
            throw new RuntimeException("获取workerId失败!");
        } finally {
            zkClient.close();
        }
    }

    private String getAddress() {
        if (DockerUtils.isDocker()) {
            return DockerUtils.getDockerHost() + DockerUtils.getDockerPort();
        } else {
            try {
                /** 同一个物理机，端口号统一为80 */
                return NetUtils.getLocalInetAddress().getHostAddress() + ":8089";
            } catch (SocketException e) {
                log.error("获取本地ip地址失败:{}", e.getMessage());
                throw new RuntimeException("获取本地ip地址失败!");
            }
        }
    }

    @Override
    public long assignDisposableWorkerId() {
        String zkAddress = zookeeperProperties.getAddress();
        String uidRoot = zookeeperProperties.getDisposableUidRoot();
        ZkClient zkClient = new ZkClient(zkAddress,
                zookeeperProperties.getSessionTimeout(), zookeeperProperties.getConnectionTimeout());
        if (!zkClient.exists(uidRoot)) {
            zkClient.createPersistent(uidRoot, true);
        }

        Stat stat = zkClient.writeDataReturnStat(uidRoot, new byte[0], -1);
        long versionAsSeq = stat.getVersion();
        zkClient.close();
        log.info("workerId={}", versionAsSeq);
        return versionAsSeq;
    }
}
